#include "ASMInstruction.hpp"
#include "CodeGen.hpp"
#include "Module.hpp"

#include <iostream>
#include <memory>
#include <unordered_map>

// 主函数声明：将高级语言翻译为RISC-V汇编
void translate_main(CodeGen *codegen);

int main() {
    // 初始化编译器和代码生成器
    auto *module = new Module();
    auto *codegen = new CodeGen(module);

    // 设置代码段起始位置
    codegen->append_inst(".text", ASMInstruction::Atrribute);

    // 执行代码翻译
    translate_main(codegen);

    // 输出生成的汇编代码
    std::cout << codegen->print();
    
    // 释放资源
    delete codegen;
    delete module;
    return 0;
}

// RISC-V汇编生成核心逻辑实现
void translate_main(CodeGen *codegen) {
    // 符号表：记录变量在栈帧中的偏移量
    std::unordered_map<std::string, int> offset_map;

    /* ================= 函数声明部分 ================= */
    codegen->append_inst(".globl main", ASMInstruction::Atrribute);  // 全局可见
    codegen->append_inst(".type main, @function", ASMInstruction::Atrribute);  // 函数类型

    /* ================= 函数入口 ================= */
    codegen->append_inst("main", ASMInstruction::Label);  // 函数标签

    /* ================= 序言部分 (Prologue) ================= */
    // 栈帧分配说明：
    // - 96字节空间满足8字节对齐要求
    // - 包含保存寄存器空间(ra, s0)和局部变量存储
    codegen->append_inst("addi sp, sp, -96");  // 栈指针调整
    codegen->append_inst("sd ra, 88(sp)");     // 保存返回地址
    codegen->append_inst("sd s0, 80(sp)");     // 保存帧指针
    codegen->append_inst("addi s0, sp, 96");   // 设置新帧指针

    /* ================= 代码生成部分 ================= */
    codegen->append_inst(".main_label_entry", ASMInstruction::Label);  // 代码段标签

    /* 数组内存分配逻辑 */
    codegen->append_inst("%op0 = alloca [10 x i32]", ASMInstruction::Comment);
    // 内存布局说明：
    // - %op0是指针变量（8字节），存储在s0-24
    // - 数组实际占用40字节（10个i32），存储在s0-64
    offset_map["%op0"] = -24;
    offset_map["array"] = -64;
    codegen->append_inst("addi", {"t0", "s0", std::to_string(offset_map["array"])});
    codegen->append_inst("sd", {"t0", "s0", std::to_string(offset_map["%op0"])});

    /* 数组首元素地址计算 */
    codegen->append_inst("%op1 = getelementptr [10 x i32], [10 x i32]* %op0, i32 0, i32 0",
                        ASMInstruction::Comment);
    offset_map["%op1"] = -32;
    codegen->append_inst("ld", {"t0", "s0", std::to_string(offset_map["%op0"])});
    codegen->append_inst("sd", {"t0", "s0", std::to_string(offset_map["%op1"])});

    /* 数组初始化操作 */
    codegen->append_inst("store i32 10, i32* %op1", ASMInstruction::Comment);
    codegen->append_inst("ld", {"t0", "s0", std::to_string(offset_map["%op1"])});
    codegen->append_inst("li t1, 10");
    codegen->append_inst("sw t1, 0(t0)");

    /* 数组第二元素地址计算 */
    codegen->append_inst("%op2 = getelementptr [10 x i32], [10 x i32]* %op0, i32 0, i32 1",
                        ASMInstruction::Comment);
    offset_map["%op2"] = -40;
    codegen->append_inst("ld", {"t0", "s0", std::to_string(offset_map["%op0"])});
    codegen->append_inst("addi", {"t0", "t0", "4"});  // i32类型偏移4字节
    codegen->append_inst("sd", {"t0", "s0", std::to_string(offset_map["%op2"])});

    /* 数据加载与运算 */
    codegen->append_inst("%op3 = load i32, i32* %op1", ASMInstruction::Comment);
    offset_map["%op3"] = -48;
    codegen->append_inst("ld", {"t0", "s0", std::to_string(offset_map["%op1"])});
    codegen->append_inst("lw", {"t1", "0(t0)"});
    codegen->append_inst("sw", {"t1", "s0", std::to_string(offset_map["%op3"])});

    /* 乘法运算实现 */
    codegen->append_inst("%op4 = mul i32 %op3, 3", ASMInstruction::Comment);
    offset_map["%op4"] = -52;
    codegen->append_inst("lw", {"t0", "s0", std::to_string(offset_map["%op3"])});
    codegen->append_inst("li", {"t1", "3"});
    codegen->append_inst("mul", {"t0", "t0", "t1"});
    codegen->append_inst("sw", {"t0", "s0", std::to_string(offset_map["%op4"])});

    /* 数据存储操作 */
    codegen->append_inst("store i32 %op4, i32* %op2", ASMInstruction::Comment);
    codegen->append_inst("lw", {"t0", "s0", std::to_string(offset_map["%op4"])});
    codegen->append_inst("ld", {"t1", "s0", std::to_string(offset_map["%op2"])});
    codegen->append_inst("sw", {"t0", "0(t1)"});

    /* 最终数据加载 */
    codegen->append_inst("%op5 = load i32, i32* %op2", ASMInstruction::Comment);
    offset_map["%op5"] = -56;
    codegen->append_inst("ld", {"t0", "s0", std::to_string(offset_map["%op2"])});
    codegen->append_inst("lw", {"t1", "0(t0)"});
    codegen->append_inst("sw", {"t1", "s0", std::to_string(offset_map["%op5"])});

    /* ================= 函数返回部分 ================= */
    codegen->append_inst("ret i32 %op5", ASMInstruction::Comment);
    codegen->append_inst("lw", {"a0", "s0", std::to_string(offset_map["%op5"])});
    // 统一退出点设计：确保资源释放逻辑集中处理
    codegen->append_inst("j main_exit");  

    /* ================= 收尾部分 (Epilogue) ================= */
    codegen->append_inst("main_exit", ASMInstruction::Label);
    codegen->append_inst("ld ra, 88(sp)");    // 恢复返回地址
    codegen->append_inst("ld s0, 80(sp)");    // 恢复帧指针
    codegen->append_inst("addi sp, sp, 96"); // 释放栈空间
    codegen->append_inst("jr ra");            // 函数返回
}
